from tree_sitter import Language, Parser
import os

os.environ['CFLAGES'] = '-std=c11'

def build_language_so():
    LANGUAGES = {
        "python": "https://githubfast.com/tree-sitter/tree-sitter-python",
        "javascript": "https://githubfast.com/tree-sitter/tree-sitter-javascript",
        "cpp": "https://githubfast.com/tree-sitter/tree-sitter-cpp",
        "java": "https://githubfast.com/tree-sitter/tree-sitter-java",
        "go": "https://githubfast.com/tree-sitter/tree-sitter-go"
    }

    if not os.path.exists("build"):
        os.mkdir("build")

    # 克隆语法库
    for name, repo in LANGUAGES.items():
        path = f"build/tree-sitter-{name}"
        if not os.path.exists(path):
            os.system(f"git clone {repo} {path}")

    Language.build_library(
        "build/my-languages.so",
        [f"build/tree-sitter-{name}" for name in LANGUAGES]
    )

def load_languages():
    lib_path = "build/my-languages.so"
    languages = {
        "python": Language(lib_path, "python"),
        "javascript": Language(lib_path, "javascript"),
        "cpp": Language(lib_path, "cpp"),
        "java": Language(lib_path, "java"),
        "go": Language(lib_path, "go"),
    }
    return languages

def chunk_code(code: str, language_name="python"):
    languages = load_languages()
    parser = Parser()
    parser.set_language(languages[language_name])

    tree = parser.parse(bytes(code, "utf8"))
    root_node = tree.root_node

    chunks = []
    for node in root_node.children:
        # 针对 Python：函数定义、类定义
        if node.type in ("function_definition", "class_definition"):
            start = node.start_point[0]
            end = node.end_point[0]
            chunk = "\n".join(code.splitlines()[start:end+1])
            chunks.append(chunk)
    return chunks

# if __name__ == '__main__':
#     LANGUAGES = {
#         "python": "https://githubfast.com/tree-sitter/tree-sitter-python",
#         "javascript": "https://githubfast.com/tree-sitter/tree-sitter-javascript",
#         "cpp": "https://githubfast.com/tree-sitter/tree-sitter-cpp",
#         "java": "https://githubfast.com/tree-sitter/tree-sitter-java",
#         "go": "https://githubfast.com/tree-sitter/tree-sitter-go"
#     }
#     Language.build_library(
#         "build/my-languages.so",
#         [f"build/tree-sitter-{name}" for name in LANGUAGES]
#     )

if __name__ == "__main__":
    code = """
        class MyClass:
            def method_one(self, x):
                return x + 1

            def method_two(self, y):
                return y * 2

        def standalone_function(z):
            return z - 3
        """
    chunks = chunk_code(code, "python")
    for i, c in enumerate(chunks, 1):
        print(f"---- Chunk {i} ----\n{c}\n")